import numpy as np
import json
from libtiff import TIFF
import glob
import os
import cv2
import nibabel as nib
from bresenham import bresenham


def mkdir(path):
    folder = os.path.exists(path)

    if not folder:
        os.makedirs(path)


def bresenham_change(first_point,second_point):
    tmp_list = list(bresenham(first_point[0],first_point[1],second_point[0],second_point[1]))
    return tmp_list


# connect all the points of an array
def plot_line(arr):
    res = []
    for i in range(len(arr) - 1):
        res.extend(bresenham_change(arr[i], arr[i + 1]))
    res.extend(bresenham_change(arr[-1], arr[0]))
    return res


def processData_3D(path, savePath):
    jsonPaths = glob.glob(path + '/*.json')

    imgs = []
    labels = []

    for i in range(len(jsonPaths)):
        patient_imgs = []
        patient_labels = []
        all_location = []
        newImgPath = []
        i_jsonPath = jsonPaths[i]
        print(i_jsonPath)
        file = open(i_jsonPath, 'rb')
        jsonfile = json.load(file)
        Esophagus = jsonfile['Esophagus']
        Heart = jsonfile['Heart']
        SpinalCord = jsonfile['Spinal-Cord']
        Lung_L = jsonfile['Lung-Left']
        Lung_R = jsonfile['Lung-Right']
        tiffEsophagus = Esophagus[2]
        tiffHeart = Heart[2]
        tiffSpinalCord = SpinalCord[2]
        tiffLung_L = Lung_L[2]
        tiffLung_R = Lung_R[2]

        (fileDir, tempFileName) = os.path.split(i_jsonPath)
        (dirName, extension1) = os.path.splitext(tempFileName)


        tiffDir = path + '/' + dirName
        tiffPaths = glob.glob(tiffDir + '/*.tiff')


        f = open('filename.txt', 'w')

        for j in range(len(tiffPaths)):
            location = 0
            imgPath = tiffPaths[j]
            f.writelines(imgPath)
            tif = TIFF.open(imgPath, mode='r')
            one_img = tif.read_image()

            one_label = np.zeros(np.shape(one_img), dtype=np.uint8)


            (imgsDir, tempImgName) = os.path.split(imgPath)
            (tiffName, extension2) = os.path.splitext(tempImgName)

            location = tiffName.split('_')[0]

            if tiffName in tiffEsophagus:
                zAndCoorEso = tiffEsophagus[tiffName]
                for i in range(len(zAndCoorEso) - 1):
                    CoorEso = zAndCoorEso[i + 1]
                    completEso = plot_line(CoorEso)
                    for k in range(len(completEso)):
                        x_coor = completEso[k][1]
                        y_coor = completEso[k][0]
                        one_label[x_coor][y_coor] = 1
                    cv2.fillPoly(one_label, [np.array(completEso)], 1)
                #
            if tiffName in tiffHeart:
                zAndCoorHea = tiffHeart[tiffName]
                for i in range(len(zAndCoorHea) - 1):
                    CoorHea = zAndCoorHea[i + 1]
                    completHea = plot_line(CoorHea)
                    for k in range(len(completHea)):
                        x_coor = completHea[k][1]
                        y_coor = completHea[k][0]
                        one_label[x_coor][y_coor] = 2

                    cv2.fillPoly(one_label, [np.array(completHea)], 2)

            if tiffName in tiffSpinalCord:
                zAndCoorSpin = tiffSpinalCord[tiffName]
                for i in range(len(zAndCoorSpin) - 1):
                    CoorSpin = zAndCoorSpin[i + 1]
                    completSpin = plot_line(CoorSpin)
                    for k in range(len(completSpin)):
                        x_coor = completSpin[k][1]
                        y_coor = completSpin[k][0]
                        one_label[x_coor][y_coor] = 3

                    cv2.fillPoly(one_label, [np.array(completSpin)], 3)

            if tiffName in tiffLung_L:
                zAndCoorLL = tiffLung_L[tiffName]
                for i in range(len(zAndCoorLL) - 1):
                    CoorLL = zAndCoorLL[i + 1]
                    completLL = plot_line(CoorLL)
                    for k in range(len(completLL)):
                        x_coor = completLL[k][1]
                        y_coor = completLL[k][0]
                        one_label[x_coor][y_coor] = 4

                    cv2.fillPoly(one_label, [np.array(completLL)], 4)
                # #
            if tiffName in tiffLung_R:
                zAndCoorLR = tiffLung_R[tiffName]
                for i in range(len(zAndCoorLR) - 1):
                    CoorLR = zAndCoorLR[i + 1]
                    completLR = plot_line(CoorLR)
                    for k in range(len(completLR)):
                        x_coor = completLR[k][1]
                        y_coor = completLR[k][0]
                        one_label[x_coor][y_coor] = 5
                    cv2.fillPoly(one_label, [np.array(completLR)], 5)
            patient_imgs.append(one_img)
            patient_labels.append(one_label)
            all_location.append(float(location))
            newImgPath.append(imgPath)

        index = list(range(len(all_location)))
        index.sort(key=lambda s: -all_location[s])

        patient_imgs = np.array(patient_imgs)[index]
        patient_labels = np.array(patient_labels)[index]



        train_x = patient_imgs.copy()
        location = np.zeros([train_x.shape[0], train_x.shape[1], train_x.shape[2]] + [3], dtype='float32')
        for m in range(train_x.shape[0]):
            for n in range(train_x.shape[1]):
                for k in range(train_x.shape[2]):
                    location[m, n, k][0] = m * (150/(train_x.shape[1]))
                    location[m, n, k][1] = n
                    location[m, n, k][2] = k
                    # print(location[m,j,k])
        train_x = train_x[:,np.newaxis,:, :]
        location = location.transpose(0, 3, 1, 2)
        new_train_x = np.concatenate((train_x, location), 1)

        for i in range(len(new_train_x)):
            print("Img: ", i)
            imgToSave = new_train_x[i]
            imgToSavePath = newImgPath[index[i]]
            (imgToSaveDir, tempImgToSaveName) = os.path.split(imgToSavePath)
            (imgToSaveName, extension3) = os.path.splitext(tempImgToSaveName)

            if not os.path.exists(savePath):
                mkdir(savePath)
            fileSavePath = savePath + '/' + dirName + '/' + imgToSaveName
            mkdir(fileSavePath)

            img = nib.Nifti1Image(imgToSave, np.eye(4))
            nib.save(img, fileSavePath + '/img.nii.gz')
            gt = nib.Nifti1Image(patient_labels[i], np.eye(4))
            nib.save(gt, fileSavePath  + '/gt.nii.gz')

        print("Image Data Size: ", np.shape(patient_imgs))

    return imgs, labels


processData_3D(os.path.abspath('..') + '/CT_Data/Training', os.path.abspath('.') + '/TrainingData2D')
processData_3D(os.path.abspath('..') + '/CT_Data/Test3', os.path.abspath('.') + '/newTestData2D_3')
